#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time : 2024/11/18 16:35
# @Author : XXX
# @Site : 
# @File : Replay_buffer.py
# @Software: PyCharm
from random import random


class Replay_Buffer:
    def __init__(self, capacity):
        self.capacity = capacity  # 经验回放的容量
        self.buffer = []  # 缓冲区
        self.position = 0

    def push(self, state_now, label_now):
        '''
        缓冲区是一个队列，容量超出时去掉开始存入的转移(transition)
        这个地方的队列规则是通过position来进行控制的，当position到达capcacity的一个周期时，
        则重置，所以不会超过capacity
        '''
        if len(self.buffer) < self.capacity:
            self.buffer.append(None)
        # self.buffer[self.position] = (state, action, reward, next_state)
        self.buffer[self.position] = (state_now, label_now)
        self.position = (self.position + 1) % self.capacity

    def sample(self, batch_size):
        mini_batch = random.sample(self.buffer, batch_size)  # 随机采出小批量转移
        state_now, label_now  = zip(*mini_batch)  # 解压成状态，动作等
        return list(state_now), list(label_now)

    def __len__(self):
        ''' 返回当前存储的量
        '''
        return len(self.buffer)